#!/usr/bin/python3
# coding=utf8
import sys

sys.path.append('/home/pi/MasterPi/')
import cv2
import time
import os
import threading
import signal
import yaml_handle
import Misc
import apriltag
import numpy as np
import pandas as pd
from ArmIK.Transform import *
from ArmIK.ArmMoveIK import *
import HiwonderSDK.Sonar as Sonar
import HiwonderSDK.Board as Board
import HiwonderSDK.mecanum as mecanum

if sys.version_info.major == 2:
    print('Please run this program with python3!')
    sys.exit(0)

AK = ArmIK()
chassis = mecanum.MecanumChassis()

range_rgb = {
    'red': (0, 0, 255),
    'blue': (255, 0, 0),
    'green': (0, 255, 0),
    'black': (0, 0, 0),
    'white': (255, 255, 255),
}

lab_data = None
HWSONAR = None

def load_config():
    global lab_data, servo_data

    lab_data = yaml_handle.get_yaml_data(yaml_handle.lab_file_path)


__target_color = ('red', 'green', 'blue')


def setTargetColor(target_color):
    global __target_color

    __target_color = target_color
    return (True, ())


# 找出面积最大的轮廓
# 参数为要比较的轮廓的列表
def getAreaMaxContour(contours):
    contour_area_temp = 0
    contour_area_max = 0
    area_max_contour = None

    for c in contours:  # 历遍所有轮廓
        contour_area_temp = math.fabs(cv2.contourArea(c))  # 计算轮廓面积
        if contour_area_temp > contour_area_max:
            contour_area_max = contour_area_temp
            if contour_area_temp > 300:  # 只有在面积大于300时，最大面积的轮廓才是有效的，以过滤干扰
                area_max_contour = c

    return area_max_contour, contour_area_max  # 返回最大的轮廓


# 夹持器夹取时闭合的角度
servo1 = 1500


# 初始位置
def initMove():
    Board.setPWMServoPulse(1, servo1, 300)
    AK.setPitchRangeMoving((0, 6, 18), 0, -90, 90, 1500)


# 设置蜂鸣器
def setBuzzer(timer):
    Board.setBuzzer(0)
    Board.setBuzzer(1)
    time.sleep(timer)
    Board.setBuzzer(0)


# 设置扩展板的RGB灯颜色使其跟要追踪的颜色一致
# def set_rgb(color):
#     if color == "red":
#         Board.RGB.setPixelColor(0, Board.PixelColor(255, 0, 0))
#         Board.RGB.setPixelColor(1, Board.PixelColor(255, 0, 0))
#         Board.RGB.show()
#     elif color == "green":
#         Board.RGB.setPixelColor(0, Board.PixelColor(0, 255, 0))
#         Board.RGB.setPixelColor(1, Board.PixelColor(0, 255, 0))
#         Board.RGB.show()
#     elif color == "blue":
#         Board.RGB.setPixelColor(0, Board.PixelColor(0, 0, 255))
#         Board.RGB.setPixelColor(1, Board.PixelColor(0, 0, 255))
#         Board.RGB.show()
#     else:
#         Board.RGB.setPixelColor(0, Board.PixelColor(0, 0, 0))
#         Board.RGB.setPixelColor(1, Board.PixelColor(0, 0, 0))
#         Board.RGB.show()

def set_rgb(color):
    if color=='None':
        Board.RGB.setPixelColor(0, Board.PixelColor(0, 0, 0))
        Board.RGB.setPixelColor(1, Board.PixelColor(0, 0, 0))
        Board.RGB.show()
    else:
        R,G,B=color
        Board.RGB.setPixelColor(0, Board.PixelColor(R, G, B))
        Board.RGB.setPixelColor(1, Board.PixelColor(R, G, B))
        Board.RGB.show()


count = 0
_stop = False
# color_list = []
get_roi = False
__isRunning = False
# detect_color = 'None'
detect_id = 'None'
start_pick_up = False
start_count_t1 = True
detect_color = [0,0,0]


# 变量重置
def reset():
    global _stop
    global count
    global get_roi
    # global color_list
    # global detect_color
    global detect_id
    global start_pick_up
    global start_count_t1
    global Threshold
    global speed
    global stopMotor
    global turn
    global forward
    global old_speed

    count = 0
    _stop = False
    get_roi = False
    detect_id = 'None'
    start_pick_up = False
    start_count_t1 = True
    speed = 40
    old_speed = 0
    Threshold = 30.0
    turn = False
    forward = True
    stopMotor = True


# app初始化调用
def init():
    print("AprialtagDetect Init")
    load_config()
    initMove()


# app开始玩法调用
def start():
    global __isRunning
    reset()
    turn = False
    forward = True
    stopMotor = True
    __isRunning = True
    print("AprialtagDetect Start")


# app停止玩法调用
def stop():
    global _stop
    global __isRunning
    _stop = True
    __isRunning = False
    set_rgb('None')
    chassis.set_velocity(0,0,0)
    print("AprialtagDetect Stop")


# app退出玩法调用
def exit():
    global _stop
    global __isRunning
    _stop = True
    __isRunning = False
    set_rgb('None')
    chassis.set_velocity(0,0,0)
    HWSONAR.setPixelColor(0, Board.PixelColor(0, 0, 0))
    HWSONAR.setPixelColor(1, Board.PixelColor(0, 0, 0))
    print("AprialtagDetect Exit")


rect = None
size_s = (160, 120)
size_m = (320, 240)
rotation_angle = 0
unreachable = False
world_X, world_Y = 0, 0
distance_data = []

def move():
    global rect
    global _stop
    global get_roi
    global __isRunning
    global unreachable
    global detect_color
    global detect_id
    global start_pick_up
    global rotation_angle
    global world_X, world_Y
    global turn
    global speed
    global __until
    global HWSONAR
    global Threshold
    global distance_data
    global stopMotor
    global forward
    global old_speed
    global playing
    
    changed = False

    print("start move")
    while playing:
        turn = False
        changed = True
    
        dist = HWSONAR.getDistance() / 10.0

        distance_data.append(dist)
        data = pd.DataFrame(distance_data)
        data_ = data.copy()
        u = data_.mean()  # 计算均值
        std = data_.std()  # 计算标准差

        data_c = data[np.abs(data - u) <= std]
        distance = data_c.mean()[0]

        if len(distance_data) == 5:
            distance_data.remove(distance_data[0])
                
        if __isRunning:
            if changed:   # 同样的速度值只设置一次
                changed = False
                chassis.set_velocity(speed,90,0)
                
            if distance <= Threshold:   # 检测是否达到距离阈值
                turn = True
                
            else:
                turn = False
                
            if detect_id != 'None' and start_pick_up:  # 检测到二维码

                print("detect id",detect_id)

                detect_color[0] = int((detect_id + 1) * 32 - 1)   #设置颜色，玩个花的
                detect_color[1] = int(((detect_id + 1) * 32 +127))%256
                detect_color[2] = int(256-(detect_id + 1) * 32)

                set_rgb(detect_color)  # 设置扩展板上的彩灯与检测到的颜色一样

                if detect_id == 1 and turn:  # 检测到红色,点头
                    print(distance)
                    setBuzzer(0.1)  # 设置蜂鸣器响0.1秒
                    
                    chassis.set_velocity(speed,90,-0.3)# 顺时针旋转
                    time.sleep(1)
                    detect_color = [0,0,0]
                    start_pick_up = False
                    set_rgb(detect_color)
                    changed = True
                    
                elif detect_id == 2 and turn:  # 检测到红色,点头
                    print(distance)
                    setBuzzer(0.1)  # 设置蜂鸣器响0.1秒
                    
                    chassis.set_velocity(speed,90,0.3)# 顺时针旋转
                    time.sleep(1)
                    detect_color = [0,0,0]
                    start_pick_up = False
                    set_rgb(detect_color)
                    changed = True

                else:  # 检测到绿色或者蓝色，则摇头
                    continue
                    for i in range(0, 3):
                        Board.setPWMServoPulse(6, 1300, 400)
                        time.sleep(0.5)
                        Board.setPWMServoPulse(6, 1700, 400)
                        time.sleep(0.5)
                        if not __isRunning:
                            continue

                    AK.setPitchRangeMoving((0, 6, 18), 0, -90, 90, 500)  # 回到初始位置
                    time.sleep(0.5)
                    detect_color = [0,0,0]
                    start_pick_up = False
                    set_rgb(detect_color)
                    
                turn = False
            else:
                time.sleep(0.01)
        else:
            if _stop:
                print('ok')
                _stop = False
                initMove()  # 回到初始位置
                chassis.set_velocity(0,0,0)
                time.sleep(1.5)
            turn = True
            forward = True
            time.sleep(0.01)

HWSONAR = Sonar.Sonar()
playing = True
# 运行子线程
th = threading.Thread(target=move)
th.setDaemon(True)
th.start()

t1 = 0
roi = ()
center_list = []
last_x, last_y = 0, 0
draw_color = range_rgb["black"]

apriltag_path = '/home/pi/.local/lib/python3.7/site-packages/libapriltag.so'
searchpath = [os.path.dirname(apriltag_path)]
detector = apriltag.Detector(searchpath=searchpath)

def run(img):

    global pub_time
    global publish_en
    global id_smallest
    global start_pick_up
    global detect_id


    img_copy = img.copy()
    img_h, img_w = img.shape[:2]
    frame_resize = cv2.resize(img_copy, size_m, interpolation=cv2.INTER_NEAREST)
    gray = cv2.cvtColor(frame_resize, cv2.COLOR_BGR2GRAY)
    detections = detector.detect(gray, return_image=False)
    tag_id = 'None'

    if len(detections) != 0:
        for i, detection in enumerate(detections):
            tag_id = int(detection.tag_id)  # 获取tag_id
            corners = np.rint(detection.corners)  # 获取四个角点
            for i in range(4):
                corners[i][0] = int(Misc.map(corners[i][0], 0, size_m[0], 0, img_w))
                corners[i][1] = int(Misc.map(corners[i][1], 0, size_m[1], 0, img_h))
            cv2.drawContours(img, [np.array(corners, np.int32)], -1, (0, 255, 255), 2)
            object_center_x = int(Misc.map(detection.center[0], 0, size_m[0], 0, img_w))
            object_center_y = int(Misc.map(detection.center[1], 0, size_m[1], 0, img_h))  # 中心点
            object_angle = int(
                math.degrees(math.atan2(corners[0][1] - corners[1][1], corners[0][0] - corners[1][0])))  # 计算旋转角
            cv2.putText(img, str(tag_id), (object_center_x - 10, object_center_y + 10), cv2.FONT_HERSHEY_SIMPLEX, 1,
                        [0, 255, 255], 2)
            # if id_smallest == 'None' or tag_id <= id_smallest:
            #     id_smallest = tag_id
            #     center_x = object_center_x
            #     center_y = object_center_y
            #     angle = object_angle
            #     data = id_smallest
            # print("tag id", int(tag_id))
            # print("center_x", object_center_x)
            # print("center_y", object_center_y)
            # print("angle", object_angle)
    if tag_id!='None':
        detect_id = tag_id
        start_pick_up=True

        # id_smallest = 'None'

    return img


def Stop(signum, frame):
    global playing

    playing = False
    print('关闭中...')
    chassis.set_velocity(0,0,0)  # 关闭所有电机

signal.signal(signal.SIGINT, Stop)

if __name__ == '__main__':
    init()
    start()
    HWSONAR = Sonar.Sonar()
    cap = cv2.VideoCapture('http://127.0.0.1:8080?action=stream')
    while playing:
        ret, img = cap.read()
        if ret:
            frame = img.copy()
            Frame = run(frame)
            frame_resize = cv2.resize(Frame, (320, 240))
            cv2.imshow('frame', frame_resize)
            key = cv2.waitKey(1)
            if key == 27:
                break
        else:
            time.sleep(0.01)
    cv2.destroyAllWindows()
    chassis.set_velocity(0,0,0)  # 关闭所有电机
    print('已关闭')

